'''
Author: LostPatrol_3
Date: 2023-11-04 17:19:48
LastEditors: LostPatrol_3
LastEditTime: 2023-11-06 19:17:07
Description: Implementation of client
'''
from PyQt5.QtWidgets import QApplication
from PyQt5.QtWidgets import QFileDialog
from PyQt5.QtWidgets import QLabel
from PyQt5.QtNetwork import QTcpSocket
from PyQt5.QtCore import pyqtSignal
from PyQt5.QtCore import pyqtSlot
from PyQt5.QtCore import QThread
from PyQt5.QtCore import QObject
from PyQt5.QtCore import QObject
from PyQt5.QtCore import QTimer
from PyQt5.QtGui import QTextImageFormat
from PyQt5.QtGui import QTextCursor
from PyQt5.QtGui import QPixmap
from PyQt5.QtGui import QImage
from deprecated import deprecated
import datetime
import json
import time
import sys
import os

from GUIFrame import MyFrame

# Transform Protocol:

# text_msg =
# {
#     'type' : 'text'
#     'data' : str
# }

# img_msg_head = 
# {
#     'type' : 'image'
#     'imageSize' : int
# }

# img_msg_body : bytes


"""
class SocketState(int):
    UnconnectedState = ... # type: QAbstractSocket.SocketState
    HostLookupState = ... # type: QAbstractSocket.SocketState
    ConnectingState = ... # type: QAbstractSocket.SocketState
    ConnectedState = ... # type: QAbstractSocket.SocketState
    BoundState = ... # type: QAbstractSocket.SocketState
    ListeningState = ... # type: QAbstractSocket.SocketState
    ClosingState = ... # type: QAbstractSocket.SocketState
"""


IP = '8.130.106.82'
PORT = 12345

PATCH_SIZE = 4096  # 4KB
MAX_FILE_LENGTH = 1024 * 1024 * 10 # 10MB
LISTEN_INTERVAL = 200

OTHER_ERR = -1
SOCKET_ERR = 0

TEXT_INVALID = 1
TEXT_SENT    = 2
TEXT_RECV    = 3

IMAGE_INVALID = 'image invalid'
IMAGE_SENT    = 'image sent'
IMAGE_RECV    = 'image recv'
IMAGE_DATA    = 'image data'


def reconnect(clientWorker):
    print("TCP disconnected. Retrying...")
    try:
        clientWorker.tcpSocket.connectToHost(IP, PORT)
        # clientWorker.tcpSocket.waitForConnected(1000)
    except Exception as e:
        print(e)
    else:
        print("TCP reconnected successfully")
    finally:
        return


# class SocketWorker(QObject):
class SendWorker(QObject):
    textMessageSent = pyqtSignal(int) 
    imageMessageSent = pyqtSignal(str)
    
    # !!!! pyqtSignal 必须为类属性时才能使用connect方法 !!!!

    def __init__(self, tcpSocket:QTcpSocket):
        super().__init__()
        self.tcpSocket = tcpSocket
        # 消息发送结果信息, 通知主线程更新UI


    @pyqtSlot(str)
    def sendTextMessage(self, text:str):
        if not self.tcpSocket or self.tcpSocket.state() == QTcpSocket.UnconnectedState:
            print(f"Abnormal socket state: {self.tcpSocket.state()}")
            self.textMessageSent.emit(SOCKET_ERR)
            reconnect(self)

        # 发送的消息不能超过4096字节
        if len(text) > PATCH_SIZE:
            self.textMessageSent.emit(TEXT_INVALID)
            return
        data = json.dumps({'type':'text', 'data':text})
        self.tcpSocket.write(data.encode())
        # self.tcpSocket.waitForBytesWritten(1000)
        self.textMessageSent.emit(TEXT_SENT)
        print(f"Text data sent successfully :{data}")


    @pyqtSlot(str)
    def sendImageMessage(self, imagePath:str):
        """
            发送一张图片分为两步
            1. 发送消息类型: image, 图片大小
            2. 发送图片的二进制数据
        """
        if not self.tcpSocket or self.tcpSocket.state() == QTcpSocket.UnconnectedState:
            print(f"Abnormal socket state: {self.tcpSocket.state()}")
            reconnect(self)

        # 发送的图片不能超过10MB
        imageSize = os.path.getsize(imagePath)
        if imageSize >= MAX_FILE_LENGTH:
            self.imageMessageSent.emit(IMAGE_INVALID)
            return
        
        data = json.dumps({'type':'image', 'imageSize': imageSize})
        self.tcpSocket.write(data.encode())
        # self.tcpSocket.waitForBytesWritten(1000)
        time.sleep(0.1)
        bytesWritten = 0
        try:
            with open(imagePath, 'rb') as f: # rb: readonly, binary
                data = f.read(MAX_FILE_LENGTH)
                bytesWritten = self.tcpSocket.write(data)
                # 发送部分可以一次发送较大的数据, 可以一次性写入缓冲区
                # self.tcpSocket.waitForBytesWritten(4000)
        except Exception as e:
            print(e)
            self.imageMessageSent.emit(str(e))
            return
        else:
            print(f"Image size (bytes):{imageSize}")
            print(f"Bytes successfully written:{bytesWritten}")
            if imageSize == bytesWritten:
                self.imageMessageSent.emit(imagePath)
                print(f"Image data sent successfully: {data}")
            else:
                print('Not all data sent successfully, operation canceled')
            


class ReceiveWorker(QObject):
    textMessageReceived = pyqtSignal(str)      # 用于发送文本消息的信号
    imageMessageReceived = pyqtSignal(bytes)   # 用于发送图片消息的信号
    
    def __init__(self, tcpSocket:QTcpSocket):
        super().__init__()
        self.tcpSocket = tcpSocket

        self.timer = QTimer()
        self.nextMsgType = 'text'
        self.imageSize = 0
        self.timer.timeout.connect(self.listenOnce)
        self.timer.start(LISTEN_INTERVAL)

    @pyqtSlot()
    def stopListening(self):
        self.timer.stop()

    @pyqtSlot()
    def startListening(self):
        self.timer.start(LISTEN_INTERVAL)


    def listenOnce(self):
        """
            Runnning in a separate thread over the lifetime of a tcp connection
        """
        self.timer.stop()
        self.timer.start(LISTEN_INTERVAL)
        if not self.tcpSocket or self.tcpSocket.state() == QTcpSocket.UnconnectedState:
            print(f"Abnormal socket state: {self.tcpSocket.state()}")
            reconnect(self)
        if self.nextMsgType == 'text': 
            self.listenText()
        elif self.nextMsgType == 'image':
            self.listenImage()


    def listenText(self):
        print("Start listening text")
        data = self.tcpSocket.read(PATCH_SIZE)
        if data is None or len(data) == 0:
            return
        print(f"Received raw data: {data}")
        try:
            data = data.decode()
            jsonData = json.loads(data)
        except Exception as e: 
            print(f"decode error: {e}")
            return
        if jsonData['type'] == 'text':
            data = jsonData['data']
            self.textMessageReceived.emit(data)
        elif jsonData['type'] == 'image':
            self.nextMsgType = 'image'
            self.imageSize = jsonData['imageSize']
        else:
            print("Unknown message type, skip this message")

    
    def listenImage(self):
        print("Start listening image")
        data = self.tcpSocket.read(MAX_FILE_LENGTH)
        if data is None or len(data) == 0:
            print("No image data received, continue listening...")
            return
        else:
            print(f'Image data received :{data}')
            print(f'Received data size: {len(data)}')
            self.nextMsgType = 'text'
            self.imageMessageReceived.emit(data)
            


class ClientGUI(MyFrame):
    # 此信号在主线程中发送, 触发子线程的发送方法
    sendMessageWorkSignal = pyqtSignal(str) 
    sendImageWorkSignal = pyqtSignal(str)
    stopListeningSignal = pyqtSignal()
    startListeningSignal = pyqtSignal()

    def __init__(self):
        super().__init__()
        
        self.sendWorker = None
        self.receiveWorker = None
        self.isListening = False
        # self.timer = QTimer()
        # self.timer.timeout.connect(self.checkSocket)
        # self.timer.start(LISTEN_INTERVAL)

        self.createFunctions()
        self.switchState(False)
        

    def switchState(self, connected:bool):
        if connected:
            self.QPushButton_applyServerIP.setEnabled(False)
            self.QPushButton_applyServerPort.setEnabled(False)
            self.QPushButton_connect.setEnabled(False)
            self.QPushButton_disconnect.setEnabled(True)
            self.QPushButton_send.setEnabled(True)
        else:
            self.QPushButton_applyServerIP.setEnabled(True)
            self.QPushButton_applyServerPort.setEnabled(True)
            self.QPushButton_connect.setEnabled(True)
            self.QPushButton_disconnect.setEnabled(False)
            self.QPushButton_send.setEnabled(False)
            

    def createFunctions(self):
        self.QLabel_userName.mousePressEvent = self.toggleEditUserName
        self.QLineEdit_userName.returnPressed.connect(self.updateUserName)

        self.QPushButton_applyServerIP.clicked.connect(self.updateServerIP)
        self.QPushButton_applyServerPort.clicked.connect(self.updateServerPort)
        self.QPushButton_send.clicked.connect(self.sendTextMessage)
        self.QPushButton_clear.clicked.connect(self.clearMessage)
        self.QPushButton_connect.clicked.connect(self.connectToServer)
        self.QPushButton_disconnect.clicked.connect(self.disconnectFromServer)

        self.QAction_file.triggered.connect(self.selectFile)
        self.QAction_picture.triggered.connect(self.selectPicture)

    def toggleEditUserName(self, event):
        # 点击用户名标签时切换为编辑模式
        self.QLabel_userName.setHidden(True)
        self.QLineEdit_userName.setHidden(False)
        self.QLineEdit_userName.setText(self.QLabel_userName.text())
        self.QLineEdit_userName.setFocus()

    @pyqtSlot()
    def updateUserName(self)->None:
        self.userName = self.QLineEdit_userName.text()
        self.QLabel_userName.setText(self.userName)
        self.QLabel_userName.setHidden(False)
        self.QLineEdit_userName.setHidden(True)

    @pyqtSlot()
    def updateServerIP(self)->None:
        self.ip = self.QLineEdit_serverIP.text()

    @pyqtSlot()
    def updateServerPort(self)->None:
        self.port = self.QLineEdit_serverPort.text()


    # @pyqtSlot()
    @deprecated(reason='cannot start/stop QTimer from another thread')
    def checkSocket(self):
        if not hasattr(self, 'tcpSocket'):
            return
        if self.tcpSocket.state() != QTcpSocket.ConnectedState:
            self.stopListeningSignal.emit()
            self.isListening = False
        elif self.tcpSocket.state() == QTcpSocket.ConnectedState and not self.isListening:
            self.startListeningSignal.emit()
            self.isListening = True
        self.timer.stop()
        self.timer.start(LISTEN_INTERVAL)
            

    @pyqtSlot()
    def connectToServer(self):
        self.tcpSocket = QTcpSocket()
        try:
            self.tcpSocket.connectToHost(self.ip, int(self.port))
        except Exception as e:
            print('Connect to server failed')
            print(e)
            return
        print('Connect to server successfully')

        self.sendThread = QThread(self)
        self.receiveThread = QThread(self)
        self.sendWorker = SendWorker(self.tcpSocket)
        self.receiveWorker = ReceiveWorker(self.tcpSocket)

        self.sendMessageWorkSignal.connect(self.sendWorker.sendTextMessage)
        self.sendImageWorkSignal.connect(self.sendWorker.sendImageMessage)
        self.stopListeningSignal.connect(self.receiveWorker.stopListening)
        self.startListeningSignal.connect(self.receiveWorker.startListening)

        self.sendWorker.textMessageSent.connect(self.addSentTextToTextEdit)
        self.sendWorker.imageMessageSent.connect(self.addSentImageToTextEdit)
        self.receiveWorker.textMessageReceived.connect(self.addReceivedTextToTextEdit)
        self.receiveWorker.imageMessageReceived.connect(self.addReceivedImageToTextEdit)
        

        self.receiveWorker.moveToThread(self.receiveThread)
        self.sendWorker.moveToThread(self.sendThread)
        self.receiveThread.start()
        self.sendThread.start()
        print('Worker thread initialized successfully')

        self.switchState(True)

    @pyqtSlot()
    def disconnectFromServer(self):
        if hasattr(self, 'tcpSocket'):
            self.tcpSocket.disconnectFromHost()
        self.sendThread.quit()
        self.receiveThread.quit()
        self.sendThread.wait()
        self.receiveThread.wait()
        self.tcpSocket = None
        self.sendWorker = None
        self.receiveWorker = None
        
        self.switchState(False)
    
    @pyqtSlot()
    def sendTextMessage(self):
        text = self.QTextEdit_inputBox.toPlainText()
        if text == '':
            return
        self.sendMessageWorkSignal.emit(text)

        
    # TODO
    def updateMemebers(self):
        pass

    def selectFile(self)->str:
        pass
    
    @pyqtSlot()
    def selectPicture(self)->str:
        # 选择图片路径
        # 使用QFileDialog.getOpenFileName()方法
        # 返回值为一个元组, 第一个元素为文件路径, 第二个元素为文件类型
        # 如果用户取消选择, 则返回空字符串
        imagePath, imageType = QFileDialog.getOpenFileName(self, 'Select picture', './', 'Image files(*.jpg *.png *.jpeg)')
        if imagePath == '':
            print('No picture selected')
            return
        self.sendImageWorkSignal.emit(imagePath)

    @pyqtSlot(str)
    def addSentImageToTextEdit(self, imagePath:str):
        cursor = self.QTextEdit_chatBox.textCursor()
        cursor.movePosition(QTextCursor.End)
        cursor.insertText('\n')
        cursor.movePosition(QTextCursor.End)
        imageFormat = QTextImageFormat()
        imageFormat.setWidth(100)
        imageFormat.setHeight(100)
        imageFormat.setName(imagePath)
        cursor.insertImage(imageFormat)
        cursor.movePosition(QTextCursor.End)
        cursor.insertText('\n')

    @pyqtSlot(bytes)
    def addReceivedImageToTextEdit(self, data:bytes):
        # 将接收到的bytes 转换成图片并显示在聊天框中
        qimage = QImage.fromData(data)
        # qimage.save(f'picture\\{datetime.date}.jpg')
        qimage.save(r'received.png')
        # pixmap = QPixmap()
        # pixmap.loadFromData(data)
        # imageLabel = QLabel()
        # imageLabel.setPixmap(pixmap)
        # cursor = self.QTextEdit_chatBox.textCursor()   
        # cursor.movePosition(QTextCursor.End)
        # cursor.insertText('\n')
        # # cursor.insertFragment(imageLabel)
        # cursor.movePosition(QTextCursor.End)
        # cursor.insertText('\n')   

    @pyqtSlot()
    def clearMessage(self):
        self.QTextEdit_chatBox.clear()
        self.QTextEdit_inputBox.clear()

    @pyqtSlot(int)
    def addSentTextToTextEdit(self, info:int):
        if info == TEXT_INVALID:
            text = '<font color="red">Message too long, failed to send!</font>'
        elif info == SOCKET_ERR:
            text = '<font color="red">Socket error, failed to send!</font>'
        elif info == TEXT_SENT:
            text = self.QTextEdit_inputBox.toPlainText()
            text = '<font color="blue">'+text+'</font>'
            self.QTextEdit_inputBox.clear()
        self.QTextEdit_chatBox.insertHtml(text+'\n')

    @pyqtSlot(str)
    def addReceivedTextToTextEdit(self, text:str):
        # self.QTextEdit_chatBox.textCursor().movePosition(QTextCursor.End)
        self.QTextEdit_chatBox.append(text)

    
    # @override
    def closeEvent(self, event):
        if hasattr(self, 'tcpSocket'):
            self.disconnectFromServer()
        event.accept()


# Socket client running on Windows
if __name__ == '__main__':
    app = QApplication(sys.argv)
    font = app.font()
    font.setFamily('Microsoft YaHei')
    app.setFont(font)

    mainWindow = ClientGUI()

    mainWindow.show()
    app.exec_()

